import unittest
import numpy as np
from geometry import Grid
from visualize import GridPlotter

import unittest
import numpy as np
import matplotlib.pyplot as plt


class TestGrid(unittest.TestCase):

    def setUp(self):
        """ 在测试之前创建多个 Grid 实例 """
        self.test_cases = [
            {"X_size": 5.0, "Y_size": 3.0, "N": 10, "origin": (2.5, 0)},  # 原点在 (2.5, 0)
            {"X_size": 6.0, "Y_size": 4.0, "N": 20, "origin": (0, 0)},  # 原点在 (0, 0)
            {"X_size": 10.0, "Y_size": 5.0, "N": 15, "origin": (-3, 1)},  # 原点在 (-3, 1)
            {"X_size": 8.0, "Y_size": 6.0, "N": 25, "origin": (4, -2)},  # 原点在 (4, -2)
        ]
        self.grids = [Grid(tc["X_size"], tc["Y_size"], tc["N"], origin=tc["origin"]) for tc in self.test_cases]

    def test_grid_properties(self):
        """ 测试网格的基本属性 """
        for i, (tc, grid) in enumerate(zip(self.test_cases, self.grids)):
            with self.subTest(f"Testing grid case {i + 1}"):
                # 检查网格的形状
                self.assertEqual(grid.X.shape, (tc["N"], tc["N"]))
                self.assertEqual(grid.Y.shape, (tc["N"], tc["N"]))

                # 计算预期的 X、Y 取值范围
                x_min_expected = 0 - tc["origin"][0]
                x_max_expected = tc["X_size"] - tc["origin"][0]
                y_min_expected = 0 - tc["origin"][1]
                y_max_expected = tc["Y_size"] - tc["origin"][1]

                # 检查 X 和 Y 坐标范围是否符合预期
                self.assertAlmostEqual(grid.X.min(), x_min_expected, places=1)
                self.assertAlmostEqual(grid.X.max(), x_max_expected, places=1)
                self.assertAlmostEqual(grid.Y.min(), y_min_expected, places=1)
                self.assertAlmostEqual(grid.Y.max(), y_max_expected, places=1)

                # 检查网格点数
                self.assertEqual(len(grid.x), tc["N"])
                self.assertEqual(len(grid.y), tc["N"])

    def test_plot_grid(self):
        """ 运行绘图测试，供人工检查 """
        for i, (tc, grid) in enumerate(zip(self.test_cases, self.grids)):
            with self.subTest(f"Plotting grid case {i + 1}"):
                print(f"Plotting test case {i + 1}: Size=({tc['X_size']}, {tc['Y_size']}), Origin={tc['origin']}")
                plotter = GridPlotter(grid)
                plotter.plot_grid()  # 运行后手动检查是否正确


# 运行测试
if __name__ == '__main__':
    unittest.main()


